/*
 *  Copyright 2008-2009 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

/*! \file diagonal.inl
 *  \brief Inline file for diagonal.h
 */

#include <cusp/array1d.h>
#include <cusp/blas.h>
#include <cusp/csr_matrix.h>
#include <cusp/detail/format_utils.h>
#include <cusp/detail/lu.h>
#include <cusp/detail/sparse_lu.h>

#include <thrust/functional.h>
#include <thrust/transform.h>

namespace cusp
{
namespace precond
{

// constructor
template <typename IndexType, typename ValueType, typename MemorySpace>
template<typename MatrixType>
ilu<IndexType,ValueType,MemorySpace>
::ilu(const MatrixType& A)
    : linear_operator<ValueType,MemorySpace>(A.num_rows, A.num_cols, A.num_rows)
{
    IndexType n = A.num_rows;
    IndexType nnz = A.num_entries;

    // float entries of ILU0 precon
    cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> A_h(A);

    // pointer of upper part of each row
    cusp::array1d<IndexType,cusp::host_memory> uptr(n);

    /* initialize iw to 0 in parallel*/
    cusp::array1d<IndexType,cusp::host_memory> iw(n, 0);

    IndexType nnzl = 0;
    /* main loop */
    for (IndexType k=0; k<n; k++)
    {
        /* working on row k */
        IndexType j1 = A_h.row_offsets[k];
        IndexType j2 = A_h.row_offsets[k+1];

        /* mark nnz in iw */
        for (IndexType j=j1; j<j2; j++)
            iw[A_h.column_indices[j]] = j;

        IndexType jrow = -1;
	IndexType j;
	IndexType jw;

        for (j=j1; j<j2; j++)
        {
            jrow = A_h.column_indices[j];

            if (jrow < k)
            {
                ValueType tl = A_h.values[j]*A_h.values[uptr[jrow]];

                A_h.values[j] = tl;

                // linear combination
                for (IndexType jj=uptr[jrow]+1; jj<A_h.row_offsets[jrow]; jj++){
                    jw = iw[A_h.column_indices[jj]];
                    A_h.values[jw] -= tl*A_h.values[jj];
		}
            }
            else
                break;
        }

        uptr[k] = j;

        if (jrow != k || A_h.values[j] == 0.0)
      		throw cusp::runtime_exception("Encountered a zero pivot in ILU factorization.");

        A_h.values[j] = 1.0 / A_h.values[j];

        /* reset iw */
        for (j=j1; j<j2; j++)
            iw[A_h.column_indices[j]] = 0;
    }

    // nnz in l part
    for (IndexType k=0; k<n; k++)
        nnzl += (uptr[k] - A_h.row_offsets[k]);

    // split precon to two CSRs, l and u
    cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> L_h(n,n,nnzl);
    cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> U_h(n,n,nnz - nnzl);

    IndexType i1 = 0;
    IndexType i2 = 0;

    L_h.row_offsets[0] = 0;
    U_h.row_offsets[0] = 0;

    for (IndexType k=0; k<n; k++)
    {
        // l part
        for (IndexType j=A_h.row_offsets[k]; j<uptr[k]; j++)
        {
            L_h.column_indices[i1] = A_h.column_indices[j];
            L_h.values[i1] = A_h.values[j];
            i1++;
        }
        L_h.row_offsets[k+1] = L_h.row_offsets[k] + uptr[k] - A_h.row_offsets[k];

        // u part
        for (IndexType j=uptr[k]; j<A_h.row_offsets[k+1]; j++)
        {
            U_h.column_indices[i2] = A_h.column_indices[j];
            U_h.values[i2] = A_h.values[j];
            i2++;
        }
        U_h.row_offsets[k+1] = U_h.row_offsets[k] + A_h.row_offsets[k+1] - uptr[k];
    }

    L = L_h;
    U = U_h;

    L_level = cusp::detail::Level<IndexType,MemorySpace,cusp::lu_forward>(L);
    U_level = cusp::detail::Level<IndexType,MemorySpace,cusp::lu_backward>(U);

    y.resize(n);
}

// linear operator
template <typename IndexType, typename ValueType, typename MemorySpace>
template <typename VectorType1, typename VectorType2>
void ilu<IndexType, ValueType, MemorySpace>
::operator()(const VectorType1& b, VectorType2& x)
{
	cusp::sparse_forwardsolve(L, L_level, b, y);	
	cusp::sparse_backsolve(U, U_level, y, x);	
}

} // end namespace precond
} // end namespace cusp

